#include "domtree.h"

#include <functional>
#include <iostream>
#include <ostream>

#include "irblock.h"

void DomTree::Build(IRBlock* entry, bool reverse) {
    nodes.clear();
    dom.clear();
    idom.clear();
    DF.clear();

    auto next = [&](IRBlock* u) {
        return reverse ? u->GetPrev() : u->GetNext();
    };
    auto prev = [&](IRBlock* u) {
        return reverse ? u->GetNext() : u->GetPrev();
    };

    std::unordered_map<IRBlock*, int> dfn;
    std::unordered_map<IRBlock*, IRBlock*> par;
    std::function<void(IRBlock*)> dfs = [&](IRBlock* x) {
        dfn[x] = nodes.size();
        nodes.push_back(x);
        for (auto y : next(x)) {
            if (dfn.count(y)) continue;
            par[y] = x, dfs(y);
        }
    };
    dfs(entry);

    std::unordered_map<IRBlock*, IRBlock*> semi;
    for (auto u : nodes) semi[u] = u;
    std::unordered_map<IRBlock*, IRBlock*> mn(semi), f(semi);

    std::function<IRBlock*(IRBlock*)> merge = [&](IRBlock* x) {
        if (x == f[x]) return x;
        auto t = f[x];
        f[x] = merge(f[x]);
        if (dfn[semi[mn[t]]] < dfn[semi[mn[x]]]) mn[x] = mn[t];
        return f[x];
    };

    int n = (int)nodes.size();
    std::unordered_map<IRBlock*, std::vector<IRBlock*>> S;
    for (int i = n - 1; i; --i) {
        int tmp = n;
        auto u = nodes[i];
        for (auto j : prev(nodes[i])) {
            if (dfn[j] < i) {
                tmp = std::min(tmp, dfn[j]);
            } else {
                merge(j);
                tmp = std::min(tmp, dfn[semi[mn[j]]]);
            }
        }
        semi[u] = nodes[tmp];
        f[u] = par[u];
        S[semi[u]].push_back(u);
        for (auto j : S[nodes[i - 1]]) {
            merge(j);
            if (semi[mn[j]] == semi[j])
                idom[j] = semi[j];
            else
                idom[j] = mn[j];
        }
    }
    for (auto u : nodes) {
        if (idom[u] != semi[u]) idom[u] = idom[idom[u]];
    }
    // for (auto u : nodes) {
    //     std::cout << u->name << ": ";
    //     for (auto v : next(u)) std::cout << v->name << ' ';
    //     std::cout << "; ";
    //     for (auto [v, w] : u->prev) std::cout << v->name << ' ';
    //     std::cout << '\n';
    // }
    for (auto u : nodes) {
        if (u == entry) continue;
        dom[idom[u]].push_back(u);
    }

    std::function<void(IRBlock*)> getDF = [&](IRBlock* u) {
        for (auto v : next(u)) {
            if (idom[v] != u) DF[u].insert(v);
        }
        for (auto v : dom[u]) {
            getDF(v);
            for (auto x : DF[v]) {
                if (idom[x] != u) DF[u].insert(x);
            }
        }
    };
    getDF(entry);
}
